import os

from langchain.llms import OpenAI
from configs import EMBEDDING_MODEL
from server.knowledge_base.migrate import create_tables, folder2db

# os.environ["SENIVERSE_KEY"] = "SCcuXrU5gsV5bPlfg"

# llm = OpenAI(model=MODEL,
#              temperature=0.9)

# prompt_template = get_prompt_template("agent_chat", "ChatGLM3")
# prompt_template_agent = CustomPromptTemplate(
#     template=prompt_template,
#     tools=tools,
#     input_variables=["input", "intermediate_steps", "history"]
# )

if __name__ == '__main__':
    os.environ["OPENAI_API_KEY"] = "EMPTY"
    os.environ["OPENAI_API_BASE"] = "http://192.168.3.3:8000/v1"
    MODEL = "chatglm3-6b"

    # from langchain.document_loaders import UnstructuredPowerPointLoader

    # pip install unstructured python-pptx
    # !pip install "unstructured[all-docs]"
    # loader = UnstructuredPowerPointLoader("knowledge_base/小学语文/1 白鹭.pptx")
    # document = loader.load()
    # print(f'documents:{len(document)}')

    # from langchain.text_splitter import RecursiveCharacterTextSplitter
    # text_splitter = RecursiveCharacterTextSplitter(
    #     chunk_size=100,
    #     chunk_overlap=20,
    #     length_function=len
    # )
    # pages = loader.load_and_split(text_splitter)
    # print(f'documents:{len(pages)}')

    # from langchain.embeddings import OpenAIEmbeddings
    # import tiktoken
    #
    # tiktoken_model_name = MODEL
    # try:
    #     encoding = tiktoken.encoding_for_model(tiktoken_model_name)
    # except KeyError:
    #     tiktoken_model_name = "text-embedding-ada-002"
    # embeddings_model = OpenAIEmbeddings(model=MODEL,
    #                                     tiktoken_model_name=tiktoken_model_name,
    #                                     default_query={"tiktoken_model_name": tiktoken_model_name})
    #
    # embeddings = embeddings_model.embed_documents([document[0].page_content])
    # print(f'documents:{len(embeddings)}')
    # create_tables()
    #  pip install faiss-cpu
    # folder2db(kb_names=[], mode="recreate_vs", embed_model=EMBEDDING_MODEL)

    query = "白鹭"
    from server.knowledge_base.kb_service.base import KBServiceFactory, SupportedVSType
    from server.knowledge_base.kb_doc_api import DocumentWithScore

    knowledge_base_name = ("小学语文", MODEL)
    kb = KBServiceFactory.get_service("小学语文", SupportedVSType.FAISS, MODEL)
    docs = kb.search_docs(query, score_threshold=0.5)
    data = [DocumentWithScore(**x[0].dict(), score=x[1]) for x in docs]

    context = "\n".join([doc.page_content for doc in data])
    print(context)
